#pragma once
#include "ChainingModeEncryption.h"
#include "ChainingModeDecryption.h"
#include "Encryptor.h"
#include "LastBlockEncryption.h"
#include "HexParser.h"
#include "OFB.h"

namespace AppliedCryptography
{
	namespace AC = AppliedCryptography;

	template<unsigned int TKeySize, unsigned int TBlockSize, class TCryptor = AC::Encryptor<TKeySize, TBlockSize>, class TLBHandle = AC::TruncatePadding<TKeySize, TBlockSize>::Encryption>
	class OFB
	{
	private:
		// as this is symetric for ancryption/decryption only one impl will be used
		class _impl : public AC::ChainingMode<TKeySize, TBlockSize, TCryptor, TLBHandle>
		{
		public:
			virtual void Process(const TCryptor& cryptor,
				const TLBHandle& lbHandler,
				const Block<TBlockSize>& iv, 
				const BlockData<TBlockSize>& input, 
				BlockData<TBlockSize>& output) const
			{
				const AC::Cryptor<TKeySize, TBlockSize>& baseCryptor = cryptor;
				const AC::LastBlockHandler<TKeySize, TBlockSize>& baseLbHandler = lbHandler;

				Block<TBlockSize> xored;
				Block<TBlockSize> chyper;
				Block<TBlockSize> feed;

				Block<TBlockSize> secondToLast;
				Block<TBlockSize> last;

				feed = iv;

				// block count should be greater than 2
				// otherwise signed/unsigned mismatch
				if(input.Blocks() > 1)
				{
					unsigned int i = 0;
					while(i < (input.Blocks() - 2))
					{
						baseCryptor.Process(feed, chyper);
						xored = input[i] ^ chyper;
						output.Add(xored, TBlockSize);
						feed = chyper;
						++i;
					};
				};

				baseCryptor.Process(feed, secondToLast);
				last = baseLbHandler.Process(baseCryptor, input[input.Blocks()-1], input.LastBlockLength(), secondToLast);

				if(input.Blocks() > 1)
				{
					secondToLast ^= input[input.Blocks() - 2];
					output.Add(secondToLast, TBlockSize);
				};

				output.Add(last, input.LastBlockLength());
			};
		};

	public:
		class Encryption : public AC::ChainingModeEncryption<TKeySize, TBlockSize, TCryptor, TLBHandle>
		{
		private:
			typename OFB<TKeySize, TBlockSize, TCryptor, TLBHandle>::_impl _ofb;
		public:
			virtual void Process(const TCryptor& cryptor,
				const TLBHandle& lbHandler,
				const Block<TBlockSize>& iv, 
				const BlockData<TBlockSize>& input, 
				BlockData<TBlockSize>& output) const
			{
				return _ofb.Process(cryptor, lbHandler, iv, input, output);
			};
		};

		class Decryption : public AC::ChainingModeDecryption<TKeySize, TBlockSize, TCryptor, TLBHandle>
		{
		private:
			typename OFB<TKeySize, TBlockSize, TCryptor, TLBHandle>::_impl _ofb;
		public:
			virtual void Process(const TCryptor& cryptor,
				const TLBHandle& lbHandler,
				const Block<TBlockSize>& iv, 
				const BlockData<TBlockSize>& input, 
				BlockData<TBlockSize>& output) const
			{
				return _ofb.Process(cryptor, lbHandler, iv, input, output);
			};
		};
	};
};
